#ifndef PCPS_RDM_HEADER
#define PCPS_RDM_HEADER

#include <src/pcps.h>
#include <src/input.h>
#include <src/wavefunction.h>
#include <src/cluster.h>
#include <src/blas_lapack.h>

namespace pcps {

  //-------------------------------------------------------------------------------
  // pcps::get_1rdm_data -- collect the 1-body RDM elements needed to find the
  //                        many body hartree-fock RDM corresponding to the
  //                        specified pairs of creation and destruction operators
  //-------------------------------------------------------------------------------

  template <class SCALAR>
  inline void get_1rdm_data(const int nsites,
                            const int nops,
                            const int * const cre_sites,
                            const int * const des_sites,
                            const SCALAR * const one_rdm,
                            SCALAR * const rdm_data)
  {
    for (int i = 0; i < nops; i++)
      for (int j = 0; j < nops; j++)
        rdm_data[i*nops+j] = one_rdm[ cre_sites[i]*nsites + des_sites[j] ];
  }

  //-------------------------------------------------------------------------------
  // pcps::compute_hf_rdm_element -- compute the many body hartree-fock RDM 
  //                                 element corresponding to the specified pairs
  //                                 of creation and destruction operators
  //-------------------------------------------------------------------------------

  template <class SCALAR>
  SCALAR compute_hf_rdm_element(const int nops,
                                const int * const occs,
                                const SCALAR * const rdm_data,
                                SCALAR * const det,
                                int * const ipiv);

  //-------------------------------------------------------------------------------
  // Compact_RDM -- class to hold the data that will be necessary to compute rdm
  //                elements efficiently while looping over configurations inside
  //                the expectation_value_detail function.
  //                This class is templated on the scalar type (SCALAR), the
  //                operator type (OP), the particle type (PART), and the reference
  //                type (REF).
  //                If there is no specialization of this class for a particular
  //                combination of SCALAR, OP, PART, and REF, then the constructor
  //                raises an error.
  //-------------------------------------------------------------------------------

  //-------------------------------------------------------------------------------
  // Compact_RDM -- raise an error if no specialization is found
  //-------------------------------------------------------------------------------

  template <class SCALAR, class OP, class PART, class REF> class Compact_RDM {

    public:

      // constructor to prepare rdm data
      Compact_RDM(const pcps::Input & userinp,
                  const pcps::Wavefunction<SCALAR> & wfn,
                  const OP & op,
                  const pcps::Cluster & local_cluster) {
        throw pcps::Exception( (boost::format("no implementation for pcps::Compact_RDM with template parameters SCALAR = %s, OP = %s, PART = %s, REF = %s")
                                 % pcps::type_lookup(typeid(SCALAR))
                                 % pcps::type_lookup(typeid(OP))
                                 % pcps::type_lookup(typeid(PART))
                                 % pcps::type_lookup(typeid(REF)) ).str() );
      }

      // function to compute an rdm element
      SCALAR element(const pcps::Cluster & local_cluster) const {
        throw pcps::Exception( (boost::format("no implementation for pcps::Compact_RDM with template parameters SCALAR = %s, OP = %s, PART = %s, REF = %s")
                                 % pcps::type_lookup(typeid(SCALAR))
                                 % pcps::type_lookup(typeid(OP))
                                 % pcps::type_lookup(typeid(PART))
                                 % pcps::type_lookup(typeid(REF)) ).str() );
        return pcps::zero<SCALAR>();
      }

  };

  //-------------------------------------------------------------------------------
  // Compact_RDM -- specialization for fermions and the RHF reference
  //-------------------------------------------------------------------------------

  template<class SCALAR, class OP> class Compact_RDM <SCALAR, OP, pcps::Part_Fermion, pcps::Ref_RHF> {
    
    // class data
    private:
      int _dim_a, _dim_b, _op_cre_site, _op_des_site;
      boost::shared_array<int> _cre_sites_a;
      boost::shared_array<int> _des_sites_a;
      boost::shared_array<int> _cre_sites_b;
      boost::shared_array<int> _des_sites_b;
      boost::shared_array<int> _occs_a;
      boost::shared_array<int> _occs_b;
      boost::shared_array<SCALAR> _a_data;
      boost::shared_array<SCALAR> _b_data;
      boost::shared_array<SCALAR> _det;
      boost::shared_array<int> _ipiv;

    public:

      // constructor to prepare rdm data
      Compact_RDM(const pcps::Input & userinp,
                  const pcps::Wavefunction<SCALAR> & wfn,
                  const OP & op,
                  const pcps::Cluster & local_cluster) {

        // create a reference to the vector of cluster sites
        const std::vector<int> & sites = local_cluster.sites();

        // determine the number of sites in the cluster
        const int n_cluster_sites = sites.size();

        // find the size of the alpha and beta determinants
        _dim_a = 0;
        _dim_b = 0;
        _op_cre_site = -1;
        _op_des_site = -1;
        if (typeid(op) == typeid(pcps::OpRL<SCALAR>)) {
          _op_cre_site = op.sites()[0];
          _op_des_site = op.sites()[1];
          if (_op_cre_site % 2 == 0 && _op_des_site % 2 == 0)
            _dim_a++;
          else if (_op_cre_site % 2 == 1 && _op_des_site % 2 == 1)
            _dim_b++;
          else
            throw pcps::Exception("expected operator sites to be both alpha or both beta");
        }
        for (int p = 0; p < n_cluster_sites; p++)
          if (sites[p] != _op_cre_site && sites[p] != _op_des_site)
            if (sites[p] % 2 == 0)
              _dim_a++;
            else
              _dim_b++;

        // prepare the lists of alpha creation/destruction site pairs
        _cre_sites_a = pcps::allocate_shared_array<int>(_dim_a);
        _des_sites_a = pcps::allocate_shared_array<int>(_dim_a);
        {
          int q = 0;
          if (typeid(op) == typeid(pcps::OpRL<SCALAR>) && _op_cre_site % 2 == 0) {
            _cre_sites_a[q] = _op_cre_site;
            _des_sites_a[q] = _op_des_site;
            q++;
          }
          for (int p = 0; p < n_cluster_sites; p++)
            if (sites[p] % 2 == 0 && sites[p] != _op_cre_site && sites[p] != _op_des_site) {
              _cre_sites_a[q] = sites[p];
              _des_sites_a[q] = sites[p];
              q++;
            }
        }

        // prepare the lists of beta creation/destruction site pairs
        _cre_sites_b = pcps::allocate_shared_array<int>(_dim_b);
        _des_sites_b = pcps::allocate_shared_array<int>(_dim_b);
        {
          int q = 0;
          if (typeid(op) == typeid(pcps::OpRL<SCALAR>) && _op_cre_site % 2 == 1) {
            _cre_sites_b[q] = _op_cre_site;
            _des_sites_b[q] = _op_des_site;
            q++;
          }
          for (int p = 0; p < n_cluster_sites; p++)
            if (sites[p] % 2 == 1 && sites[p] != _op_cre_site && sites[p] != _op_des_site) {
              _cre_sites_b[q] = sites[p];
              _des_sites_b[q] = sites[p];
              q++;
            }
        }

        // allocate arrays to hold alpha and beta site occupations
        _occs_a = pcps::allocate_shared_array<int>(_dim_a);
        _occs_b = pcps::allocate_shared_array<int>(_dim_b);
        
        // allocate arrays to hold the 1-body rdm data
        _a_data = pcps::allocate_shared_array<SCALAR>(_dim_a * _dim_a);
        _b_data = pcps::allocate_shared_array<SCALAR>(_dim_b * _dim_b);

        // allocate arrays for use in computing the determinant via the LU decomposition
        const int max_dim = std::max(_dim_a, _dim_b);
        _det = pcps::allocate_shared_array<SCALAR>(max_dim * max_dim);
        _ipiv = pcps::allocate_shared_array<int>(max_dim);

        // collect the 1-body alpha rdm data
        pcps::get_1rdm_data(userinp.nsites(), _dim_a, _cre_sites_a.get(), _des_sites_a.get(), &wfn.hf_rdm()[0], _a_data.get());

        // collect the 1-body beta rdm data
        pcps::get_1rdm_data(userinp.nsites(), _dim_b, _cre_sites_b.get(), _des_sites_b.get(), &wfn.hf_rdm()[0], _b_data.get());

      }

      // function to compute an rdm element
      SCALAR element(const pcps::Cluster & local_cluster) {

        // initialize the return value
        SCALAR retval = pcps::unity<SCALAR>();

        // get occupations of alpha sites
        for (int p = 0; p < _dim_a; p++)
          if (_cre_sites_a[p] == _des_sites_a[p])
            _occs_a[p] = local_cluster.ket_occ(_cre_sites_a[p]);
          else
            _occs_a[p] = 1;

        // compute the alpha determinant
        retval *= pcps::compute_hf_rdm_element(_dim_a, _occs_a.get(), _a_data.get(), _det.get(), _ipiv.get());

        // get occupations of beta sites
        for (int p = 0; p < _dim_b; p++)
          if (_cre_sites_b[p] == _des_sites_b[p])
            _occs_b[p] = local_cluster.ket_occ(_cre_sites_b[p]);
          else
            _occs_b[p] = 1;

        // compute the beta determinant
        retval *= pcps::compute_hf_rdm_element(_dim_b, _occs_b.get(), _b_data.get(), _det.get(), _ipiv.get());

        // return the rdm element
        return retval;

      }

  };

  //-------------------------------------------------------------------------------
  // Compact_RDM -- specialization for spinless fermions and the uniform reference
  //-------------------------------------------------------------------------------

  template<class SCALAR, class OP> class Compact_RDM <SCALAR, OP, pcps::Part_SpinlessFermion, pcps::Ref_Uniform> {
    
    // class data
    private:
      int _n_evn_clust_sites;
      int _n_odd_clust_sites;
      int _n_evn_outer_sites;
      int _n_odd_outer_sites;
      std::vector<int> _evn_clust_sites;
      std::vector<int> _odd_clust_sites;
      std::vector<SCALAR> _data;

    public:

      // constructor to prepare rdm data
      Compact_RDM(const pcps::Input & userinp,
                  const pcps::Wavefunction<SCALAR> & wfn,
                  const OP & op,
                  const pcps::Cluster & local_cluster) {

        const int n_total_sites = userinp.nsites();
        const int n_total_parts = userinp.nparticles();

        const std::set<int> cluster_sites(local_cluster.sites().begin(), local_cluster.sites().end());

        // separate the cluster sites by parity
        for (std::set<int>::const_iterator site = cluster_sites.begin(); site != cluster_sites.end(); site++)
          if (typeid(OP) == typeid(pcps::OpRL<SCALAR>))
            if ( (*site < op.sites()[0] && *site > op.sites()[1]) || (*site > op.sites()[0] && *site < op.sites()[1]) )
              _odd_clust_sites.push_back(*site);
            else
              _evn_clust_sites.push_back(*site);
          else
            _evn_clust_sites.push_back(*site);

        // count the number of cluster sites with even and odd parity
        _n_evn_clust_sites = _evn_clust_sites.size();
        _n_odd_clust_sites = _odd_clust_sites.size();

        // count the number of outer sites with even and odd parity
        _n_evn_outer_sites = 0;
        _n_odd_outer_sites = 0;
        for (int site = 0; site < n_total_sites; site++)
          if (cluster_sites.count(site) == 0)
            if (typeid(OP) == typeid(pcps::OpRL<SCALAR>))
              if ( (site < op.sites()[0] && site > op.sites()[1]) || (site > op.sites()[0] && site < op.sites()[1]) )
                _n_odd_outer_sites++;
              else
                _n_evn_outer_sites++;
            else
              _n_evn_outer_sites++;

        // precompute the possible rdm elements
        _data.assign((_n_evn_clust_sites+1)*(_n_odd_clust_sites+1), pcps::zero<SCALAR>());
        const double n_total_confg = pcps::n_choose_m(n_total_sites, n_total_parts);
        for (int n_odd_clust_parts = 0; n_odd_clust_parts <= _n_odd_clust_sites; n_odd_clust_parts++)
          for (int n_evn_clust_parts = 0; n_evn_clust_parts <= _n_evn_clust_sites; n_evn_clust_parts++) {
            // get a reference to the rdm element
            const int local_cmpd_index = (_n_evn_clust_sites+1)*n_odd_clust_parts + n_evn_clust_parts;
            SCALAR & value = _data.at(local_cmpd_index);
            // determine the number of particles inside and outside the cluster
            const int n_clust_parts = n_evn_clust_parts + n_odd_clust_parts;
            const int n_outer_parts = n_total_parts - n_clust_parts;
            // compute the rdm element
            for (int n_odd_outer_parts = 0; n_odd_outer_parts <= n_outer_parts; n_odd_outer_parts++) {
              const int n_evn_outer_parts = n_outer_parts - n_odd_outer_parts;
              value += pcps::unity<SCALAR>()
                       * std::pow(-1.0, n_odd_clust_parts + n_odd_outer_parts)
                       * pcps::n_choose_m(_n_odd_outer_sites, n_odd_outer_parts)
                       * pcps::n_choose_m(_n_evn_outer_sites, n_evn_outer_parts);
            }
            value /= n_total_confg;
          }

      }

      // function to compute an rdm element
      SCALAR element(const pcps::Cluster & local_cluster) const {

        // compute number of even particles in the cluster
        int n_evn_clust_parts = 0;
        for (int i = 0; i < _n_evn_clust_sites; i++)
          n_evn_clust_parts += local_cluster.ket_occ(_evn_clust_sites[i]);

        // compute number of odd particles in the cluster
        int n_odd_clust_parts = 0;
        for (int i = 0; i < _n_odd_clust_sites; i++)
          n_odd_clust_parts += local_cluster.ket_occ(_odd_clust_sites[i]);

        // return the corresponding rdm element
        const SCALAR retval = _data[(_n_evn_clust_sites+1)*n_odd_clust_parts + n_evn_clust_parts];
        //std::cout << "cluster config = ";
        //for (int i = 0; i < local_cluster.sites().size(); i++)
        //  std::cout << boost::format("%2i ") % local_cluster.ket_occ(local_cluster.sites().at(i));
        //std::cout << boost::format("  rdm element = %s") % pcps::format_number("%.6f", retval) << std::endl;
        return retval;

      }

  };

  //-------------------------------------------------------------------------------
  // Compact_RDM -- specialization for spinless fermions and the RHF reference
  //-------------------------------------------------------------------------------

  template<class SCALAR, class OP> class Compact_RDM <SCALAR, OP, pcps::Part_SpinlessFermion, pcps::Ref_RHF> {
    
    // class data
    private:
      int _dim, _op_cre_site, _op_des_site;
      boost::shared_array<int> _cre_sites;
      boost::shared_array<int> _des_sites;
      boost::shared_array<int> _occs;
      boost::shared_array<SCALAR> _data;
      boost::shared_array<SCALAR> _det;
      boost::shared_array<int> _ipiv;

    public:

      // constructor to prepare rdm data
      Compact_RDM(const pcps::Input & userinp,
                  const pcps::Wavefunction<SCALAR> & wfn,
                  const OP & op,
                  const pcps::Cluster & local_cluster) {

        // create a reference to the vector of cluster sites
        const std::vector<int> & sites = local_cluster.sites();

        // determine the number of sites in the cluster
        const int n_cluster_sites = sites.size();

        // find the size of the determinant
        _dim = 0;
        _op_cre_site = -1;
        _op_des_site = -1;
        if (typeid(op) == typeid(pcps::OpRL<SCALAR>)) {
          _op_cre_site = op.sites()[0];
          _op_des_site = op.sites()[1];
          _dim++;
        }
        for (int p = 0; p < n_cluster_sites; p++)
          if (sites[p] != _op_cre_site && sites[p] != _op_des_site)
            _dim++;

        // prepare the lists of creation/destruction site pairs
        _cre_sites = pcps::allocate_shared_array<int>(_dim);
        _des_sites = pcps::allocate_shared_array<int>(_dim);
        {
          int q = 0;
          if (typeid(op) == typeid(pcps::OpRL<SCALAR>)) {
            _cre_sites[q] = _op_cre_site;
            _des_sites[q] = _op_des_site;
            q++;
          }
          for (int p = 0; p < n_cluster_sites; p++)
            if (sites[p] != _op_cre_site && sites[p] != _op_des_site) {
              _cre_sites[q] = sites[p];
              _des_sites[q] = sites[p];
              q++;
            }
        }

        // allocate an array to hold the site occupations
        _occs = pcps::allocate_shared_array<int>(_dim);
        
        // allocate an array to hold the 1-body rdm data
        _data = pcps::allocate_shared_array<SCALAR>(_dim * _dim);

        // allocate arrays for use in computing the determinant via the LU decomposition
        _det = pcps::allocate_shared_array<SCALAR>(_dim * _dim);
        _ipiv = pcps::allocate_shared_array<int>(_dim);

        // collect the 1-body rdm data
        pcps::get_1rdm_data(userinp.nsites(), _dim, _cre_sites.get(), _des_sites.get(), &wfn.hf_rdm()[0], _data.get());

      }

      // function to compute an rdm element
      SCALAR element(const pcps::Cluster & local_cluster) {

        // get the site occupations
        for (int p = 0; p < _dim; p++)
          if (_cre_sites[p] == _des_sites[p])
            _occs[p] = local_cluster.ket_occ(_cre_sites[p]);
          else
            _occs[p] = 1;

        // compute and return the determinant
        return pcps::compute_hf_rdm_element(_dim, _occs.get(), _data.get(), _det.get(), _ipiv.get());

      }

  };

  //-------------------------------------------------------------------------------
  // Compact_RDM -- specialization for spins and the uniform reference
  //-------------------------------------------------------------------------------

  template<class SCALAR, class OP> class Compact_RDM <SCALAR, OP, pcps::Part_Spin, pcps::Ref_Uniform> {
    
    // class data
    private:
      std::vector<double> _rdm;

    public:

      // constructor to prepare rdm data
      Compact_RDM(const pcps::Input & userinp,
                  const pcps::Wavefunction<SCALAR> & wfn,
                  const OP & op,
                  const pcps::Cluster & local_cluster)
        : _rdm(userinp.configs_per_site() * local_cluster.size(), 0.0)
      {

        // get the number of configurations on a site
        const int configs_per_site = userinp.configs_per_site();

        // get the number of sites in the lattice
        const int n_lattice_sites = userinp.nsites();

        // z spin on one site = (occ of site) - (configs per site - 1) / 2.0

        // total z spin in lattice = (total occ of lattice) - (# sites in lattice) * (configs per site - 1) / 2.0

        // determine the required total occupation of the lattice
        const double temp = userinp.z_spin() + double(n_lattice_sites * (configs_per_site-1)) / 2.0 + 0.01;
        //if (temp != std::floor(temp))
        //  throw pcps::Exception( ( boost::format("impossible combination of particle spin (%.2f) and total z spin (%.2f)")
        //                           % userinp.particle_spin() % userinp.z_spin() ).str() );
        const int required_lattice_occ = int(temp);

        // get the total number of lattice configurations in the uniform reference
        const double total_lattice_configs = pcps::n_integer_solutions(n_lattice_sites, required_lattice_occ, configs_per_site-1);

        // determine the number of sites outside the cluster
        const int N = userinp.nsites() - local_cluster.size();

        // for each total occupation of the cluster, compute the rdm element
        for (int cluster_occ = 0; cluster_occ < _rdm.size(); cluster_occ++) {

          // get the total occupation of the outer sites
          const int M = required_lattice_occ - cluster_occ;

          // compue the rdm element
          _rdm[cluster_occ] = pcps::n_integer_solutions(N, M, configs_per_site-1) / total_lattice_configs;

        }

      }

      // function to compute an rdm element
      SCALAR element(const pcps::Cluster & local_cluster) const {

        // get the cluster's total occupation
        const int cluster_occ = local_cluster.total_occ();
        //for (int i = 0; i < local_cluster.size(); i++)
        //  cluster_occ += local_cluster.allowed_occs()[i][local_cluster.occ_nums()[i]];

        // return the rdm value
        return pcps::unity<SCALAR>() * _rdm[cluster_occ];

      }

      // function to compute an rdm element given the total occupation
      SCALAR element(const int cluster_occ) const { return pcps::unity<SCALAR>() * _rdm[cluster_occ]; }

  };

}

#endif
